#!/bin/bash
# Script which restores metadata into a VDI
# Citrix Systems Inc, 2008

if [ ! -e @INVENTORY@ ]; then
  echo Must run on a Citrix Hypervisor host.
  exit 1
fi

. @INVENTORY@

XE="@BINDIR@/xe"

master_uuid=$(${XE} pool-list params=master --minimal)
if [ $? -gt 0 ]; then
   echo Error: unable to determine master host uuid
   exit 1
fi

if [ "${master_uuid}" != "${INSTALLATION_UUID}" ]; then
   echo Error: must run this script on the master host in the resource pool.
   exit 1
fi

metadata_version=1
default_restore_mode="all"
debug="/bin/true"

function usage {
    echo "Usage: $0 [-h] [-v] [-y] [-n] [-p] [-f] [-x <VDI UUID>] [-u <SR UUID>] [-m all|sr]"
    echo
    echo " -h: Display this help message"
    echo " -x: Specify the VDI UUID to override probing"
    echo " -p: Just scan for a metadata VDI and print out its UUID to stdout"
    echo " -u: UUID of the SR you wish to restore from"
    echo " -n: Perform a dry run of the metadata import commands (default: false)"
    echo " -l: Just list the available backup dates"
    echo " -d: Specify which date to restore from (default: latest)"
    echo " -m: Either 'sr' to restore only the VMs on the SR, or 'all' for VMs (default: ${default_restore_mode})"
    echo " -v: Verbose output"
    echo " -y: Assume non-interactive mode and yes to all questions"
    echo " -f: Forcibly restore VM metadata, dangerous due to its destructive nature, please always do a dry run before using this (default: false)"
    echo 
    exit 1
}

function test_sr {
  sr_uuid_found=$(${XE} sr-list uuid="$1" --minimal)
  if [ "${sr_uuid_found}" != "$1" ]; then
     echo Invalid SR UUID specified: $1
     usage
  fi
}

dry_run=0
sr_uuid=
yes=0
just_list_dates=0
just_probe=0
chosen_date=""
restore_mode=${default_restore_mode}
force=0
while getopts "yhpvx:d:lnu:m:f" opt ; do
    case $opt in
    h) usage ;;
    u) sr_uuid=${OPTARG} ;;
    n) dry_run=1 ;;
    l) just_list_dates=1 ;;
    p) just_probe=1 ;;
    v) debug="" ;;
    d) chosen_date=${OPTARG} ;;
    m) restore_mode=${OPTARG} ;;
    x) vdis=${OPTARG} ;;
    y) yes=1 ;;
    f) force=1 ;;
    *) echo "Invalid option"; usage ;;
    esac
done

case "${restore_mode}" in
all) ;;
sr) ;;
*) echo Unknown restore mode: "${restore_mode}"; usage
esac

# get pool uuid
IFS=,
pool_uuid=$(${XE} pool-list params=uuid --minimal)
if [ -z "${pool_uuid}" ]; then
  echo Unable to determine pool UUID.
  exit 1
fi

# determine if the SR UUID is vaid
if [  -z "${sr_uuid}" ]; then
  # use the default-SR from the pool
  sr_uuid=$(${XE} pool-param-get uuid=${pool_uuid} param-name=default-SR)
fi
test_sr "${sr_uuid}"

sr_name=$(${XE} sr-param-get uuid=${sr_uuid} param-name=name-label)

# get a list of all VDIs if an override has not been provided on the cmd line
if [ -z "${vdis}" ]; then
  vdis=$(${XE} vdi-list params=uuid sr-uuid=${sr_uuid} read-only=false --minimal)
fi

mnt=
function cleanup {
   cd /
   if [ ! -z "${mnt}" ]; then
      umount ${mnt} >/dev/null 2>&1
      rmdir ${mnt}
      mnt=""
   fi

   if [ ! -z "${vbd_uuid}" ]; then
      ${debug} echo -n "Unplugging VBD: " >&2
      ${XE} vbd-unplug uuid=${vbd_uuid} timeout=20
      # poll for the device to go away if we know its name
      if [ "${device}" != "" ]; then
          device_gone=0
          for ((i=0; i<10; i++)); do
              ${debug} echo -n "." >&2
              if [ ! -b ${device} ]; then
                  ${debug} echo " done" >&2
                  device_gone=1
                  break
              fi
              sleep 1
          done
          if [ ${device_gone} -eq 0 ]; then
              ${debug} echo " failed" >&2
              ${debug} echo Please destroy VBD ${vbd_uuid} manually. >&2
          else
              ${XE} vbd-destroy uuid=${vbd_uuid}
              vbd_uuid=""
          fi
      fi
      device=""
   fi
}

if [ -z "${vdis}" ]; then
   echo No VDIs found on SR. >&2
   exit 0
fi

trap cleanup SIGINT ERR

for vdi_uuid in ${vdis}; do
   ${debug} echo -n "Creating VBD: " >&2
   vbd_uuid=$(${XE} vbd-create vm-uuid=${CONTROL_DOMAIN_UUID} vdi-uuid=${vdi_uuid} device=autodetect 2>/dev/null)

   if [ $? -ne 0 -o -z "${vbd_uuid}" ]; then
      ${debug} echo error creating VBD for VDI ${vdi_uuid} >&2
      cleanup
      continue
   fi

   ${debug} echo ${vbd_uuid} >&2

   ${debug} echo -n "Plugging VBD: " >&2
   ${XE} vbd-plug uuid=${vbd_uuid}
   device=/dev/$(${XE} vbd-param-get uuid=${vbd_uuid} param-name=device)

   if [ ! -b ${device} ]; then
     ${debug} echo ${device}: not a block special >&2
     cleanup
     continue
   fi

   ${debug} echo ${device} >&2

   ${debug} echo -n "Probing device: " >&2
   probecmd="@LIBEXECDIR@/probe-device-for-file"
   metadata_stamp="/.ctxs-metadata-backup"
   mnt=
   ${probecmd} ${device} ${metadata_stamp}
   if [ $? -eq 0 ]; then
     ${debug} echo found metadata backup >&2
     ${debug} echo -n "Mounting filesystem: " >&2
     mnt=/var/run/pool-backup-${vdi_uuid}
     mkdir -p ${mnt}
     /sbin/fsck -a ${device} >/dev/null 2>&1
     if [ $? -ne 0 ]; then
        echo File system integrity error.  Please correct manually. >&2
        cleanup
        continue
     fi
     mount ${device} ${mnt} >/dev/null 2>&1
     if [ $? -ne 0 ]; then
       ${debug} echo failed >&2
       cleanup
     else
       if [ -e "${mnt}/.ctxs-metadata-backup" ]; then
          ${debug} echo Found backup metadata on VDI: ${vdi_uuid} >&2
          xe vdi-param-set uuid=${vdi_uuid} other-config:ctxs-pool-backup=true
          break
       fi
     fi
   else
     ${debug} echo backup metadata not found >&2
   fi
   cleanup
done

if [ $just_probe -gt 0 ]; then
   echo ${vdi_uuid}
   cleanup
   exit 0
fi

cd ${mnt}
${debug} echo "" >&2

if [ ! -d ${mnt}/metadata ]; then
   echo Metadata backups not found. >&2
   cleanup
   exit 1
fi

cd ${mnt}/metadata

if [ $just_list_dates -gt 0 ]; then
    ls -1r ${mnt}/metadata
    cleanup
    exit 0
fi

if [ -z "${chosen_date}" ]; then
    chosen_metadata_dir=$(ls | sort -n | tail -1)
    if [ -z "${chosen_metadata_dir}" ]; then
       echo No VM metadata backups found in ${mnt}/metadata >&2
       cleanup
       exit 1
    fi
else
    if [ ! -d ${mnt}/metadata/${chosen_date} ]; then
       echo Date directory "${chosen_date}" not found >&2
       cleanup
       exit 1
    fi
    chosen_metadata_dir=${chosen_date}
fi

case ${restore_mode} in
sr)
    full_dir=${mnt}/metadata/${chosen_metadata_dir}/by-sr/${sr_uuid}
    ;;
all)
    full_dir=${mnt}/metadata/${chosen_metadata_dir}/all
    ;;
esac

if [ ! -d ${full_dir} ]; then
    echo No VM metadata exports were found for the selected SR >&2
    cleanup
    exit 1
fi

${debug} echo Selected: ${full_dir}

cd ${full_dir}
${debug} echo "" >&2
${debug} echo Latest VM metadata found is: >&2
${debug} ls >&2

if [ $yes -eq 0 ]; then
   echo "Do you wish to reimport all VM metadata?" 
   echo "Please type in "yes" and <enter> to continue." 
   read response
   if [ "$response" != "yes" ]; then
     echo Aborting metadata restore.
     cleanup
     exit 1
   fi
fi

${debug} echo "" >&2
${debug} echo Restoring VM metadata: >&2

trap - ERR

IFS=" "
error_count=0
good_count=0
if [ ${force} -gt 0 ]; then
    force_flag=" --force"
else
    force_flag=""
fi
if [ ${dry_run} -gt 0 ]; then
    dry_run_flag=" --dry-run"
else
    dry_run_flag=""
fi
shopt -s nullglob
for meta in *.vmmeta; do
   echo xe vm-import filename=${meta} sr-uuid=${sr_uuid} --metadata --preserve${force_flag}${dry_run_flag}
   "@BINDIR@/xe" vm-import filename="${full_dir}/${meta}" sr-uuid=${sr_uuid} --metadata --preserve${force_flag}${dry_run_flag}
   if [ $? -gt 0 ]; then
      error_count=$(( $error_count + 1 ))
   else
      good_count=$(( $good_count + 1 ))
   fi
done

smmeta_file=${mnt}/metadata/${chosen_metadata_dir}/SRMETA.xml
if [ "$restore_mode" == "all" ]; then
   cmd="@LIBEXECDIR@/restore-sr-metadata.py -f ${smmeta_file}"
else
   cmd="@LIBEXECDIR@/restore-sr-metadata.py -u ${sr_uuid} -f ${smmeta_file}"
fi

if [ -e ${smmeta_file} ]; then
    if [ ${dry_run} -gt 0 ]; then
        echo ${cmd}
    else
        ${cmd}
    fi
fi

echo "Restored ${good_count} VM(s), and ${error_count} error(s)"
cleanup
exit ${error_count}
